import torch
import torch.nn
import onnx
from torchvision import transforms
import torch.nn as nn
from torch.nn import Sequential

from model import ResNet18, Residual


# 判断是否GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# device1 = torch.device('cpu')
# 定义网络
model = ResNet18(Residual)
model.load_state_dict(torch.load('best_model.pth',map_location='cuda:0'))
model.eval()

input_names = ['input']
output_names = ['output']

# x = torch.randn(1,3,32,32,requires_grad=True)
x = torch.randn(1, 3, 224, 224, requires_grad=True)  # 这个要与你的训练模型网络输入一致。我的是黑白图像

torch.onnx.export(model, x, 'ResNet18-1.onnx', input_names=input_names, output_names=output_names, verbose='True')
